import numpy as np
from astropy.time import Time
from astropy.utils import iers
import astropy.units as u


def sec2rad(deg=0.0, minute=0.0, sec=0.0):
    return np.pi * (deg * 3600 + minute * 60 + sec) / (180 * 3600)


def rotation_matrix(axis, theta):
    if axis == 'x': return np.array([[1, 0, 0],
                                     [0, np.cos(theta), np.sin(theta)],
                                     [0, -np.sin(theta), np.cos(theta)]
                                     ])

    if axis == 'y': return np.array([[np.cos(theta), 0, -np.sin(theta)],
                                     [0, 1, 0],
                                     [np.sin(theta), 0, np.cos(theta)]
                                     ])

    if axis == 'z': return np.array([[np.cos(theta), np.sin(theta), 0],
                                     [-np.sin(theta), np.cos(theta), 0],
                                     [0, 0, 1]
                                     ])


class CoordinatesTransformEquinoxBase:
    def __init__(self, utc_time: Time):
        self.time = utc_time
        self.delta_mu = 0.0
        self.nutation_matrix()
        # self.mid_matrix = mid_matrix
        # self.matrix_list = list()

    def precession_matrix(self):

        T = (self.time.tt.mjd - 51544.5) / 36525.0

        xi = 2306.2181 * T + 0.30188 * (T ** 2) + 0.017998 * (T ** 3)
        z = 2306.2181 * T + 1.09468 * (T ** 2) + 0.018203 * (T ** 3)
        theta = 2004.3109 * T - 0.42665 * (T ** 2) - 0.041833 * (T ** 3)

        xi = sec2rad(sec=xi)
        z = sec2rad(sec=z)
        theta = sec2rad(sec=theta)

        return rotation_matrix('z', -z) @ rotation_matrix('y', theta) @ rotation_matrix('z', -xi)

    def nutation_matrix(self):
        T = (self.time.tt.mjd - 51544.5) / 36525.0

        params = [[0, 0, 0, 0, 1, -171996, -174.2, 92025, 8.9],
                  [0, 0, 2, -2, 2, -13187, -1.6, 5736, -3.1],
                  [0, 0, 2, 0, 2, -2274, -0.2, 977, -0.5],
                  [0, 0, 0, 0, 2, 2062, 0.2, -895, 0.5],
                  [0, 1, 0, 0, 0, 1426, -3.4, 54, -0.1],
                  [1, 0, 0, 0, 0, 712, 0.1, -7, 0.0],
                  [0, 1, 2, -2, 2, -517, 1.2, 224, -0.6],
                  [0, 0, 2, 0, 1, -386, -0.4, 200, 0.0],
                  [1, 0, 2, 0, 2, -301, 0.0, 129, -0.1],
                  [0, -1, 2, -2, 2, 217, -0.5, -95, 0.3],
                  [1, 0, 0, -2, 0, -158, 0.0, -1, 0.0],
                  [0, 0, 2, -2, 1, 129, 0.1, -70, 0.0],
                  [-1, 0, 2, 0, 2, 123, 0.0, -53, 0.0],
                  [1, 0, 0, 0, 1, 63, 0.1, -33, 0.0],
                  [0, 0, 0, 2, 0, 63, 0.0, -2, 0.0],
                  [-1, 0, 2, 2, 2, -59, 0.0, 26, 0.0],
                  [-1, 0, 0, 0, 1, -58, -0.1, 32, 0.0],
                  [1, 0, 2, 0, 1, -51, 0.0, 27, 0.0],
                  [2, 0, 0, -2, 0, 48, 0.0, 1, 0.0],
                  [-2, 0, 2, 0, 1, 46, 0.0, -24, 0.0]]

        alpha_with_T = np.array([
            sec2rad(134, 57, 46.733) + (sec2rad(deg=1325 * 360) + sec2rad(198, 52, 2.633)) * T + sec2rad(sec=31.310) * (
                    T ** 2),
            sec2rad(357, 31, 39.804) + (sec2rad(deg=99 * 360) + sec2rad(359, 3, 1.224)) * T - sec2rad(sec=0.577) * (
                    T ** 2),
            sec2rad(93, 16, 18.877) + (sec2rad(deg=1342 * 360) + sec2rad(82, 1, 3.137)) * T - sec2rad(sec=13.257) * (
                    T ** 2),
            sec2rad(297, 51, 1.307) + (sec2rad(deg=1236 * 360) + sec2rad(307, 6, 41.328)) * T - sec2rad(sec=6.891) * (
                    T ** 2),
            sec2rad(125, 2, 40.280) - (sec2rad(deg=5 * 360) + sec2rad(134, 8, 10.539)) * T + sec2rad(sec=7.455) * (
                    T ** 2)
        ])

        delta_psi = 0.0
        delta_epsilon = 0.0
        for i in range(20):
            angle = 0
            for j in range(5):
                angle = angle + params[i][j] * alpha_with_T[j]
            delta_psi = delta_psi + (sec2rad(sec=params[i][5] * 1e-4) + sec2rad(sec=params[i][6] * 1e-4) * T) * np.sin(
                angle)
            delta_epsilon = delta_epsilon + (
                    sec2rad(sec=params[i][7] * 1e-4) + sec2rad(sec=params[i][8] * 1e-4) * T) * np.cos(angle)

        epsilon_A = sec2rad(deg=23, minute=26, sec=21.448) - sec2rad(sec=46.815) * T
        delta_mu = delta_psi * np.cos(epsilon_A)
        self.delta_mu = delta_mu
        delta_theta = delta_psi * np.sin(epsilon_A)

        return rotation_matrix('x', -delta_epsilon) @ rotation_matrix('y', delta_theta) @ rotation_matrix('z', -delta_mu)

    def earth_rotation_matrix(self):
        t = (self.time.ut1.mjd - 51544.5) / 36525.0
        S_G_hat = 18.6973746 + 879000.0513367 * t + (0.093104 * (t ** 2) - 6.2e-6 * (t ** 3)) / 3600
        S_G_hat_theta = S_G_hat * 2 * np.pi / 24.0
        # sg = self.time.sidereal_time('apparent','greenwich').to('rad').value
        return rotation_matrix('z', S_G_hat_theta + self.delta_mu)

    def polar_motion_matrix(self):
        iers_table = iers.IERS_Auto.open()
        x_p, y_p = iers_table.pm_xy(self.time)
        # x_p = iers_table.pm_x(utc_time)
        # y_p = iers_table.pm_y(utc_time)
        return rotation_matrix('y', -x_p) @ rotation_matrix('x', -y_p)


def read_txt(file_path):
    data_list = list()
    with open('data.txt') as f:
        lines = f.readlines()
        for line in lines:
            # print(line)
            str1, str2, str3 = line.split(' ')
            data_list.append([float(str1) * u.s, float(str2) * u.deg, float(str3) * u.deg])
    return data_list

def newton_generation(initial_value, e_1, M_1, tolerance):
    """
    牛顿迭代求解开普勒方程 E - e sin E = M
    :param initial_value: 开始迭代的值，可以为任意值，下面代码中取了上一个时刻的E为初始值
    :param e_1: 加了_1后缀是为了和下面代码中的e和M进行区分
    :param M_1: 同上
    :param tolerance: 迭代到误差小于什么时候停止，下面代码中取了1e-6
    :return: 计算结果
    """
    difference = 2 * tolerance
    old_value = initial_value
    new_value = initial_value
    while difference > tolerance:
        new_value = old_value - (old_value - e_1 * np.sin(old_value) * u.rad - M_1) / (1 - e_1 * np.cos(old_value))
        difference = abs(old_value - new_value)
        old_value = new_value

    return new_value


def get_orbit_elements(r_0, r_dot_0):
    # 下面计算轨道根数，先得到目标的初始时刻位置和速度
    object_initial_location = r_0 * u.m
    object_initial_speed = r_dot_0 * u.m / u.s

    # 由目标初始位置和速度得到h
    h_vector = np.cross(object_initial_location, object_initial_speed)
    h_scalar = np.linalg.norm(h_vector)

    # 由h矢量得到方向向量R
    R_vector = h_vector / h_scalar

    # 再由方向向量得到i和Ω
    i = np.arccos(R_vector[-1])
    omega = np.arcsin(R_vector[0] / np.sin(i))

    # 引入μ，参考书本附录，模型是WGS84, 因为测站的坐标转换中，第32行代码里from_geodetic，默认选择了WGS84模型
    mu = 398600.4418 * 1e9 * (u.m ** 3) / ((u.s) ** 2)

    # 计算目标初始时刻位置和速度的大小
    r_scalar = np.linalg.norm(object_initial_location)
    v_scalar = np.linalg.norm(object_initial_speed)

    # 计算轨道半长轴
    a = 1 / (2 / r_scalar - v_scalar ** 2 / mu)

    # 计算E, e, M
    E = np.arctan2(np.dot(object_initial_speed, object_initial_location) / np.sqrt(mu * a), (1 - r_scalar / a))
    e = (1 - r_scalar / a) / np.cos(E)
    M = E - e * np.sin(E) * u.rad

    # 计算向量P和Q，用于后续得到目标位置和速度的矢量
    P_vector = np.cos(E) / r_scalar * object_initial_location - np.sqrt(a / mu) * np.sin(E) * object_initial_speed
    Q_vector = (np.sin(E) / r_scalar * object_initial_location + np.sqrt(a / mu) * (
                np.cos(E) - e) * object_initial_speed)

    # 计算ω， 但后续没用到
    w = np.arctan2(P_vector[-1], Q_vector[-1])

    return [a, e, i, omega, w, M]

